import jieba
import random


class EntityExtractor:
    """
    实体抽取器: 负责将文本中的实体信息提取出来
    """

    def __init__(
            self,
            focus_on: set,
            name2entities: dict,
            ignore_entity: set = None,
            stop_word: set = None
    ):
        """
        :param focus_on: 该提取器中会进行AIML代码替换的实体类型
        :param name2entities: 该提取器用到的实体名字到对应实体列表(因为可能有多个)的映射
        :param ignore_entity: 该提取器中会无视掉(替换成'')的实体
        :param stop_word: 提取器会跳过的词
        """
        self.focus_on = focus_on
        self.ignore_entity = ignore_entity
        self.stop_word = stop_word
        self.name2entities = name2entities

    def extract(self, string):
        """
        根据计划进行实体替换
        :param string:
        :return: 替换后的字符串:str, 存在的实体字典:dict
        """
        entity_dict = {}
        # 分词
        for word in jieba.cut(string, cut_all=False):
            if self.stop_word and word in self.stop_word:
                # 忽略词跳过
                continue
            if word not in self.name2entities.keys():
                # 如果不属于aiml实体, 跳过
                continue
            # 找出实体
            entity = random.choice(self.name2entities[word])

            # 如果在实体提取器的忽略实体中则替换成''
            if self.ignore_entity and entity in self.ignore_entity:
                string = string.replace(word, '')
                continue

            if entity.entity_type not in self.focus_on:
                # 如果不在这个方案的替换计划中跳过
                continue
            # 将找出的实体加入返回值
            entity_type_name = entity.entity_type.type_name
            if entity.entity_type not in entity_dict.keys():
                entity_dict[entity_type_name] = [entity.proper_name]
            else:
                entity_dict[entity_type_name].append(entity.proper_name)
            # 将出现的所有实体词全部替换成对应的类型的aiml代码
            string = string.replace(word, entity.entity_type.aiml_code)
        return string, entity_dict
